#!/usr/bin/env python3
# This file is part of Xpra.
# Copyright (C) 2011-2024 Antoine Martin <antoine@xpra.org>
# Copyright (C) 2008, 2009, 2010 Nathaniel Smith <njs@pobox.com>
# Xpra is released under the terms of the GNU GPL v2, or, at your option, any
# later version. See the file COPYING for details.

import unittest
import binascii
from time import monotonic
from xpra.util.str_fn import hexstr, strtobytes
from xpra.util.env import envbool

from xpra.net.crypto import (
    DEFAULT_SALT, DEFAULT_ITERATIONS, DEFAULT_KEYSIZE, DEFAULT_KEY_HASH, DEFAULT_IV,
    PADDING_PKCS7, pad,
    crypto_backend_init,
)

SHOW_PERF = envbool("XPRA_SHOW_PERF", False)


def log(_message) -> None:
    # print(message[:256])
    pass


def all_identical(*items: bytes) -> None:
    first = items[0]
    size = len(items)
    for i in range(size):
        assert items[i] == first


class TestCrypto(unittest.TestCase):

    def test_crypto(self) -> None:
        from xpra.net.crypto import get_modes
        for mode in get_modes():
            self.do_test_roundtrip(mode=mode)
            self.do_test_fixed_output(mode=mode)

    def do_test_fixed_output(self, mode) -> None:
        for padding in (False, True):
            enc_data = self.do_test_roundtrip(mode=mode, padding=padding)[0]
            expected = {
                "CBC": (
                    "6339a0861384e208ed20313a64912298",
                    "6339a0861384e208ed20313a649122980c79458182b8882dbf1cca137a0f88bd",
                ),
                "GCM": (
                    "6643bd873e794c3ab30a84cd18ddd5aa",
                    "6643bd873e794c3ab30a84cd18ddd5aa558654e6345cbd03240d080e736eaff3",
                ),
                "CFB": (
                    "221373bd31b628311b3f99fc6ab1fa12",
                    "221373bd31b628311b3f99fc6ab1fa127f1e6c8ac17681cbfdc96e4fe060e020",
                ),
                "CTR": (
                    "221373bd31b628311b3f99fc6ab1fa12",
                    "221373bd31b628311b3f99fc6ab1fa1244703739a928ff644c221e3f59c51807",
                ),
            }.get(mode)
            if not expected:
                print(f"warning: no fixed output test data recorded for {mode!r}")
                print(f" got {hexstr(enc_data)}")
                return
            ref_data = expected[int(padding)]
            if hexstr(enc_data) != ref_data:
                raise RuntimeError(f"expected {ref_data} but got {hexstr(enc_data)} for {mode!r} with {padding=}")

    def do_test_roundtrip(self, message=b"some message1234",
                          encrypt_count=1,
                          decrypt_count=1,
                          mode="CBC",
                          padding=True) -> list[bytes]:
        from xpra.net.crypto import get_cipher, get_key

        key_data = b"this is our secret"
        key_salt = DEFAULT_SALT
        key_hash = DEFAULT_KEY_HASH
        iterations = DEFAULT_ITERATIONS
        block_size = DEFAULT_KEYSIZE
        # test key stretching:
        args = key_data, key_salt, key_hash, block_size, iterations
        secret = get_key(*args)
        log("%s%s=%s" % (get_key, args, hexstr(secret)))
        assert secret is not None
        # test creation of encryptors and decryptors:
        iv = DEFAULT_IV
        args = secret, strtobytes(iv), mode
        enc = get_cipher(*args).encryptor()
        log("%s%s.encryptor()=%s" % (get_cipher, args, enc))
        assert enc is not None
        dec = get_cipher(*args).decryptor()
        log("%s%s.decryptor()=%s" % (get_cipher, args, dec))
        assert dec is not None
        # test encoding of a message:
        encrypted: list[bytes] = []
        for i in range(encrypt_count):
            data = message
            if padding:
                count = 33 - (len(data) + 1) % 32
                data += pad(PADDING_PKCS7, count)
            v = enc.update(data) + enc.finalize()
            assert v is not None
            if i < 10:
                encrypted.append(v)
        assert encrypted
        # test decoding of the message:
        decrypted: list[bytes] = []
        for i in range(decrypt_count):
            v = dec.update(encrypted[i % len(encrypted)]) + dec.finalize()
            log("%s%s=%s" % (dec.update, (encrypted[0],), hexstr(v)))
            assert v is not None
            if i < 10:
                decrypted.append(v)
        if decrypted:
            all_identical([message] + decrypted)
        return encrypted

    def do_test_perf(self, size=1024 * 4, enc_iterations=20, dec_iterations=20) -> list[float]:
        asize = (size + 15) // 16
        times = []
        data = b"0123456789ABCDEF" * asize
        start = monotonic()
        self.do_test_roundtrip(data, enc_iterations, dec_iterations)
        end = monotonic()
        elapsed = max(0.0001, end - start)
        speed = (asize * 16) * (enc_iterations + dec_iterations) / elapsed
        iter_time = elapsed * 1000 / (enc_iterations + dec_iterations)
        print("%10iKB: %5.1fms: %16iMB/s" % (asize * 16 // 1024, iter_time, speed // 1024 // 1024))
        times.append(end - start)
        return times

    def test_perf(self) -> None:
        if not SHOW_PERF:
            return
        RANGE = (1024, 1024 * 1024)
        print("Python Cryptography:")
        print("  Encryption:")
        for i in RANGE:
            self.do_test_perf(i, 10, 0)
        print("  Decryption:")
        for i in RANGE:
            self.do_test_perf(i, 1, 10)
        print("  Overall:")
        for i in RANGE:
            self.do_test_perf(i, 10, 10)

    def test_cbc_roundtrip(self):
        samples = {
            "bd0b0000f42ec28568656c6c6f3c86646967657374cb83786f728b686d61632b7368613531328d686d61632b736861335f3531328d686d61632b736861335f3338348d280058335f3235360e003532348b380016380c00fa303235368b686d61632b7368613232348c686d61632b626c616b6532738c686d61632b626c616b6532628b73616c742d646967657374ca8b686d61632b7368618100443531328d4d000f8f004df6ff8c636f6d70726573736f7273c3836c7a348662726f746c69846e6f6e6588656e636f64657273c28b72656e636f6465706c7573846e6f6e6585666c75736843836c7a346780438662726f746c69678043846e6f6e656780438b72656e636f6465706c75736780438770616464696e6767876f7074696f6e73c186504b43532337856d6f64657367876f7074696f6e73c4834342438347434d8343464283435452877374726574636867876f7074696f6e73c18650424b44463293707974686f6e2d63727970746f6772617068796880438776657273696f6ec3290007846d6f64658c583131207365616d6c6573738c73657373696f6e2d74797065887365616d6c657373847575696436343a306636623537383938303662363763363861363332643237653931396633396163373536383335616165623630646265313434333738346336393666323763318a6d616368696e652d6964a065636238333837356261323234653133393666653763306136613062383263378776657273696f6e83362e328a73746172745f74696d654066fe21e98c63757272656e1200fd23f48c656c61707365645f74696d650a8b7365727665725f747970658e507974686f6e2f67746b2f7831318b7365727665722efa00f20b88686f73746e616d65866665646f72618f726561646f6e6c792d4d002443881100ff07448a7365727665722d6c6f67808a6d616368696e655fd30010045d01f15c5f6e616d6585787465726d8a6b65795f726570656174c20000946b65795f7265706561745f6d6f64696669657273438462656c6c4487637572736f7273448c6465736b746f705f73697a65c23f04a43f059b887864672d6d656e75668b737562636f6d6d616e6473f080853901118d0600f3ff9a2d6465736b746f708d73746172742d6d6f6e69746f7286736861646f778d736861646f772d73637265656e87757067726164658f757067726164652d6465736b746f708f757067726164652d6d6f6e69746f728e757067726164652d736861646f77866174746163688873657373696f6e73886c61756e63686572836775698973746172742d6775698a6275672d7265706f727487746f6f6c626f788561626f7574876578616d706c658664657461636884696e666f8269648776657273696f6e8473746f7084657869748a73637265656e73686f7487636f6e74726f6c857072696e74857368656c6c8973656e642d66696c6589636f6e66696775726585636c65616e8d636c65616e2d736f636b6574738e636c65616e2d646973706c617973896175746f73746172748a73686f77636f6e6669678973657475702d73736c8873686f772d73736c846c6973748d6c6973742d73657373696f6e738c6c6973742d77696e646f777388646973706c61797388656e636f64696e6789706174682d696e666f8568746d6c3584646f6373896c6973742d6d646e73886d646e732d677569906d61785fa70103df01631e003f10e087a300f2fb823a318f636c69656e742d73687574646f776e438773686172696e67448e73686172696e672d746f67676c6544846c6f636b448b6c6f636b2d746f67676c65448777696e646f777343886b6579626f617264438d696e7075742d6465766963657385787465737490706f696e7465722e72656c617469766543876e6574776f726b678f62616e6477696474682d6c696d697400857368656c6c439273746172742d6e65772d636f6d6d616e64734492657869742d776974682d6368696c6472656e44977365727665722d636f6d6d616e64732d7369676e616c73c686534947494e54875349475445524d86534947485550875349474b494c4c87534947555352318753494755535232949403064700f216696e666f4393637572736f722e64656661756c745f73697a65208f637572736f722e6d61785301c23e403e408d726573697a655fef02f3ff7a438c726573697a655f6578616374438c666f7263655f756e6772616243976b6579626f6172642e666173742d737769746368696e67438d776865656c2e707265636973654490706f696e7465722e6f7074696f6e616c438f746f7563687061642d646576696365448c73637265656e2d73697a6573e7c23f1e003f10e0c23f14003f0b40c23f10003f0900c23f0f003f0870c23f0c803f0708c23f0b403f0654c23f0a003f0640c23f0a003f05a0c23f08003f0600c23f07803f05a0c23f07403f0570c23f07003f0540c23f08003f0480c23f07803f04b0c23f07803f0438c23f06403f04b0c23f06903f041ac23f05783f041ac23f06403f0384c23f05003f0400c23f05783f0384c23f05003f03c0c23f05583f0300c23f05003f0320c23f04803f0360c23f05003f02d0c23f04003f0300c23f04003f0240c23f03403f0270c23f03c03f021cc23f03203f0258c23f03603f01e6c23f02803f01e0c23f02d03f0195c23f02d03f0190c23f02803f0190c23f02803f0168c23f02803f015ec23f04a43f059b874201fbb767856772616273438677696e646f776b8d6672616d652d657874656e7473438f636f6e6669677572652e64656c746143877369676e616c73c686534947494e54875349475445524d875349475155495487534947434f4e548753494755535231875349475553523289647261676e64726f704386737461746573c88969636f6e696669656487666f63757365648a66756c6c73637265656e8561626f76658562656c6f7786737469636b798969636f6e6966696564896d6178696d697a65649361637475616cc803f8020f003f087090726f6f745f77696e646f771800f3168a656e6372797074696f6e728663697068657283414553846d6f6465834342438c6d6f646559020e0208978269769030303030300500c4886b65795f73616c7431362f1400040200011c00f134686173688453484131886b65795f73697a65208b6b65795f737472657463688650424b444632936b65795f737472657463682e6f7074696f6e73c18650424b444632966300027d08f30f5f697465726174696f6e733f27108770616464696e6786504b435323378f0f0004c10013c118001491e306f40c5f6b6579636f6465736e857368696674c2c28753686966745f4c000a00f2115200846c6f636bc1c289436170735f4c6f636b0087636f6e74726f6cc2c289430a00625f4c00c289431106f43e5f5200846d6f6431c4c2864d6574615f4c01c285416c745f4c01c200864d6574615f4cc2864d6574615f5201846d6f6432c1c2884e756d5f4c6f636b00846d6f6433c3c28753757065725f4c000a0015520a00f30a4c01846d6f6434c2c28748797065725f4c01c28748797065727800fd0835c2c29049534f5f4c6576656c335f536869667400c2001400f02a9270696e672d6563686f2d736f7572636569644388656e636f64696e67846175746f926175746f5f726566726573685f64656c61793f01f47f":
                "9080ed5f000fc0efea6148046553745f748a5c83c5732c3bb082ce5e479daeecca56755deec748b443104738ddb928e4005895baf85a0e1f19880ef80524d6c34ff375a4bd2f74a06aae1402a6dfde453d5a45e02893ed2ce808e24c2fa9599e80ffd4fdce1a14603857ced7f38d8f6e85c89ea1e81638c4eed906ebce67ff3f3976ea691e9f35fa528035307fecd4a4a0f38bcfd82eb8fff0d312f762c05b404763e6f91f63ff2c78f4fc8beb6c2a5cceec7a215375ba4638192c401faedea0fe0adce266b8c19fe1f05a4cd435cf17569fcb397693f1e2a5987e71548364ad48cf73a4c31bbee59962d6331d761eabdcb060029af5b98e6c57e2770516cac66a5666e14e9c7e8140f763ec732688d5334557ec2e70e3dc939595938def061ad2f62dcef1964ebc6e436c62596c41cc4e0f83295d3cac2d3f3202eaee1e63d4e49c2527e32bd085a35a63f11f0d8999bb9df67d0aafe1430f3125e0a406507ed55db6207c21e554dc25c446b3da7692b15b790fd355928b62f389ff2cbbdb706aaf6e78b4a4ca575a0fb73b959d5771b017ab070cba9381be091b0fd4bd8deacb57509a2f299ea52fbe74f5ee35e8d66a382a7d15e3c1894877db17d644d29c2ba1a373fe3ae6a515c443295ca4bcbb46f63f42972eefdcf238cbe3436a313f7eec0fe5887cb7ed652b020f595a29e1b67931440166456720a3c1e836b54433e526fdfaee09a70ba8a6f4922e6cff9273f47ecbbea31cb5a7d241072ad32c72178a5713b12768602f90e127820643d3853f8d6b351ba94346f56ad4cf3b36c5b6bff3982f52c7155b64ae50eed30e47f7262f40ae7066e5dbd94630c97e8ff7c418cf40218c18978d5d29775b23fff56fa3d25acd424e89722283880c9dd88a547af275d59cc8935d0e4d59b181413d140ce9ce6a653ca913808b122d21e9bd838035076315476cb78289112d34f8172234f17c301ece9193e2c026e19a29a70b74b0ed4b880d8a7ec7ba937289b864925ff55b1d27d8742fa2eeef7416bb627b7534a5e1d7bc6a654109f400b191aaa8c1e2ddccfa95a3c9d80d01cee23de95a33d0f18e474d9680ea6c87166196ec5813e1b325e531954842d6c7e794eea813b3bd2a6bbb2d0dcc01cd41e7bfb57cd664e8bace03e6faa21721bcb7499d8a279f22a4ec655b9219cfd2584b9081888ae5796973404af30bd54d83bc33636befb52d0ebc8457d254ef3e7ebdaa4aeb6032ef1a3c5b437b06071b6ba1e4080dd2d2fc3fc7a834048e620f4b3833a538cf6f09d077e798398cb55c3fb46fbc3f642ac25b58fddaf1fc51a6496f1effde24e782d56e9b358396fed3925f33ef1ca5aa5e3fc1661b31c195eb1006543c53dd4a7a1f6146ac32045d9f6416f2d65c8a670f1c57c3eee3d305ae0ba9a0052bc9b0e8f4c48a821115519e3286e85b6daeed5b6d434ede450c952b55dac7b8c49f3734c38caf1ce1b44c4b46b92a2010097690e59c0b47db608bb9c61af9a82af81eaafdd25ff44a71176422ae22307d9c08833142408e5f531caeab8619e7862c2d8fb2b6bd5bf33c4b350f9cf6f4b83d12e42aaf68e669259f6b9f6b05a46b15ebecfb760fe3ebba01615714817b8913eea9fa14a6419e1dbc56094eb84ecc88f0ce85c2434bce666ad046fb70ed689132489f1cb651c1495938593d5fb03103eab2bbc1f9861c0e354d9f3787c4d5253774490d2309bf2ba064b0213602dde232149e2dea0f31fd937e7672760d7448970ac985a471aa2e0ed4066185e2265baf716f3ca58ffe5ababfb0826a167d12c50736cfc59e577b7d2092e448ad70a1e4081773eaca13497c74b3bd968aa6206ca7136f564ec2a3aa45aae29d42a83d37a83bf5901b463ff114093cdbf9552decf5f6af2d9e8dc2eba6426bdf221f0485fdad27e73e7f0ca58813d263fe0e9cb0427a04fb309659d2de463a2224381d9f44f62494f6f851b22cd30f57ce26324162e79989d8dde611c622eeb0ab09978db0ff5bfd23bf3f39dcd053fb7ee2250373ca23df1c6e345d280b567c4753ba378249df0eca60ea5a2564b4c49cff56c568a4417e33e1a9a543b0fffb3c92fcb204085d3c9911f0aeb5f0b163f4c05c57ad321fa291e02541f9d170da6574649de56e33cb6670700c1a1c173cdac4918e7176f30871165d700aedf380bb1016dddfa9e334e574b004b81f4ad6b8974ebc26aa1784a523c11d258373ceb5a5183e92843960861d917d1da1b66628969956f6012648695aec15f81cd311cf164d05c11eec03799019bd61975341280d7899f267e9d304f9d558e8a23737ccdd922a18f59413bace67137deee8e540c5b330eb89d66cfa3cf90effd8951e567b8771309fbc6f0061a1949c9b5a329a5a33fbb9ca2d4da721c57ffad659360acb5aa44c8480e556fc03d81826bbf207a4ca5c80b8769fe8628decfed6d1b15264bd262a8bd2cc8bf007324a5417c9bf33e0a886c34e03b192800cfc57796ca6756bfd92631f40abe9819d38a2af2b235ee2bf94264b9f56cd37e4d910422dbcfc39c4e672e0aa858f34e9b9d77cd3a666fe45c285a0cdce5f4b450f1e950d4f110babbb9ce24fe7ad67bcf38cbe13efed3057b2bf90bf2144f0bbe5895d1606670d673f0493e29466b053563ba91090191bcab799c4de818778cdb777e6e388748e2369365e4e384fe653bc26c50221ffcbf61eae4c016489e9248b573b7ae8e78b0bdf0b88d472429a4b55f177a638c00bb67891e67dc88e3377049264f6a34181f60a923006f8ad87f85a654e6084b2ce65eb5b3d8faec06a27194a48aee85817596ae2e35c087826927f70b896059d85a58290c865aa3d43ef1db33d8b87ac5705636af303600a5c9afb5edb7bcda1c7c78e9ce538e567a94ada819f6b9289191cf2c641247acd2b169e1f4206bb7f19bd9e2264f28b8d3ed44b40ae5d0366487f6285514c944f533ab23638c0d3afb5437ee29ad686319f68452cc2f67779e390eb488d561df7711030b6357d250f4b3cd9c25d43e2f16e892340a50746dd895e42117ea8a383326b073262c7ceb99523ebffa4261712f038b9b800e265024899d416ab33391c27e54468431a2f56df635169dd5484835fcba14d56dc8d6d862e610a7cca5ff1d9a37a2746f4f94e97e2861785c66ef6a4337fe88b421579148172e10648f8e18916ad5bdcca19ae27a00ab91b989246f38816e3b76a8c63e5ae6d72e7ec69854bff721dd375474c1cd472a99f043703a45f8e9624130f1f2395d41724388ce6cad2bd82203c285e0bf360707a6ef17dfd240f2e8ee53d376ed414b7c8de9a5ff9daf22e962c57fa6e7acef60b4f795081d477316d7b1738898751c16dde321ac7bed38d262dcf4a1f6e70c51241ce7e5c4e80f50869995ffc1653cca11afa258aa704fbd632ad767d339edfb32359e19a907229f37f4c2fd8bea9ac2acd22b4ef5d0cf3aa662244a3721aa63e307c10f96e93e4334743e69b13bcdcc091ea7cfff2f471b62782f5d93e6868956cf3c64dd4524a02b3d44292e30e64753ee77e0c39b4f244fed9b5d50bfb2f8e4c9c17639e9ee9b738f62a9e7e5fc6567f1995dbbe6d0a3cfce04291c389c5e45f62cfde2569b63c46b96f27510b979d10d795ffe4eb8c8267fa66515e0d60a6082b9fce81710837e7e241d75b3097712806f604df3795015e978af603388001dfb47858224f1bb7777109e9b9c3a105cc05c246bfde0193158c9c0fe57c44b15fd8078c486689a7d4362f35853b142810d3789833e7513cc6b8c5",
            "cb84647261770100003f01f33f013c846a706567302f020068877175616c6974793e6485666c75736800":
                "d0affb4fa5dca339db3aeea4975009fb5cd95f2720e1c293af844e1c9e67cda535281d8fc476328e498b767d5704d228540b7df61a1b59fe27865801963d92f2",
            "c79277696e646f772d6d6f76652d726573697a65013f01593f022e3f01f33f013c00":
                "ec347c53d9396ae179d3c8f9c846de99273a9fc939522d3b10665d895600112109d1a3eee432c7024754bc407af97705cedc3d8c767a83ceef5caa16fda6c054",
        }
        for pdata, penc in samples.items():
            message = binascii.unhexlify(pdata)
            enc = self.do_test_roundtrip(message, mode="CBC", padding=True)[0]
            if penc != hexstr(enc):
                raise RuntimeError("expected %s but got %s" % (penc, hexstr(enc)))

    def setUp(self):
        assert crypto_backend_init(), "failed to initialize python-cryptography"


def main():
    unittest.main()


if __name__ == '__main__':
    main()
